#include "mesh.hpp"
#include <limits>
#include <iostream>
#include <fstream>
#include <sstream>
#include <cmath>
#include <Eigen/Dense>

using Eigen::MatrixXd;

istream& operator>>(istream& is, Point& p) {
	for (size_t i = 0; i < 3; ++i)
		is >> p.coords[i];
	return is;
}

ostream& operator<<(ostream& os, const Point& p) {
	os << "(" << p.coords[0] << ", " << p.coords[1] << ", " << p.coords[2] << ")";
	return os;
}

void HexMesh::initialize() {
	for (auto& elem : elems) {
		auto hexElem = dynamic_cast<HexElement*>(elem.get());
		// find the bounding box of the element
		hexElem->coords_lo.fill(std::numeric_limits<double>::max());
		hexElem->coords_hi.fill(std::numeric_limits<double>::lowest());
		for (size_t i : hexElem->nodes_id) {
			for (size_t j = 0; j < 3; ++j) {
				hexElem->coords_lo[j] = nodes[i].coords[j] < hexElem->coords_lo[j] ? nodes[i].coords[j] : hexElem->coords_lo[j];
				hexElem->coords_hi[j] = nodes[i].coords[j] > hexElem->coords_hi[j] ? nodes[i].coords[j] : hexElem->coords_hi[j];
			}
		}
        
        // compute hex centroid
        Eigen::Vector3d hexCentroid(0, 0, 0);
        for (size_t i : hexElem->nodes_id) {
            hexCentroid += Eigen::Vector3d(nodes[i].coords[0], nodes[i].coords[1], nodes[i].coords[2]);
        }
        hexCentroid /= hexElem->nodes_id.size();

		// compute normal vectors
		for (int i = 0; i < 6; ++i) { // loop over faces
        /*
			size_t id_0 = hexElem->nodes_id[HexElement::faces[i][0]]; 
			size_t id_1 = hexElem->nodes_id[HexElement::faces[i][1]];
			size_t id_3 = hexElem->nodes_id[HexElement::faces[i][3]];
			Point p0 = { nodes[id_0] }, p1 = { nodes[id_1] }, p3 = { nodes[id_3] };
			hexElem->normals_inward[i] = Vector3D(p0, p1).cross(Vector3D(p0, p3));
        */
            array<size_t, 4> face_key;
            // using all four points
            Eigen::Matrix<double, 4, 3> points;
            for (int j = 0; j < 4; ++j) {
                size_t id = hexElem->nodes_id[HexElement::faces[i][j]];
                face_key[j] = id;
                points.row(j) = Eigen::Vector3d(nodes[id].coords[0], nodes[id].coords[1], nodes[id].coords[2]);
            }
            // 对节点ID排序以确保相同面具有相同的key
            sort(face_key.begin(), face_key.end());

            Eigen::Vector3d normal;
            // Compute the centroid of the points
            Eigen::Vector3d faceCentroid = points.colwise().mean();
            hexElem->faceCentroids[i] = faceCentroid; // 存储面质心

            // 检查是否已经计算过这个面的法向量
            auto it = face_normals_cache.find(face_key);
            if (it != face_normals_cache.end()) {
                // 如果找到缓存的法向量，需要确保方向朝内
                normal = it->second;
            } else {
                // Center the points around the centroid
                Eigen::Matrix<double, 4, 3> centeredPoints = points.rowwise() - faceCentroid.transpose();

                // Perform SVD and obtain the normal as the singular vector corresponding to the smallest singular value
                Eigen::JacobiSVD<Eigen::MatrixXd> svd(centeredPoints, Eigen::ComputeThinU | Eigen::ComputeThinV);
                normal = svd.matrixV().col(2);
            }

            // determine inward normal
            Eigen::Vector3d toCentroid = hexCentroid - faceCentroid;
            if (normal.dot(toCentroid) < 0) {
                normal = -normal;
            }

            hexElem->normals_inward[i] = Vector3D(normal[0], normal[1], normal[2]);
            // 缓存计算得到的法向量
            face_normals_cache[face_key] = normal;
		}
	}
}

bool HexMesh::is_in_elem(const Point& p, const HexElement& elem) const
{
	for (size_t i = 0; i < 6; ++i) {
		Eigen::Vector3d faceCentroid = elem.faceCentroids[i];
        Vector3D v = Vector3D({ faceCentroid[0], faceCentroid[1], faceCentroid[2] }, p);
		if (v.inner(elem.normals_inward[i]) < -TOLERANCE) // 添加容差
			return false;
	}
	return true;
}

int HexMesh::find_containing_cell(const Point& p) const
{
	for (size_t iel = 0; iel < elems.size(); ++iel) { // loop over elements in the mesh
		auto hexElem = dynamic_cast<const HexElement*>(elems[iel].get());
		if (!hexElem) {
			cerr << "Element type not supported!" << endl;
			return -2;
		}
		if (is_in_elem(p, *hexElem))
			return iel;
	}
	return -1; // indicating that we didn't find one
}

double HexMesh::interpolate(const Point& p, const string& s) const
{
	int idx = find_containing_cell(p);
	if (idx == -1) {
		cerr << "Point not in mesh!" << endl;
		return std::nan("");
	}
	else {
		return elems[idx]->data[s];
	}
}

Vector3D Vector3D::cross(const Vector3D& other) const {
	return Vector3D(
		y * other.z - z * other.y,  // Calculate the x component
		z * other.x - x * other.z,  // Calculate the y component
		x * other.y - y * other.x   // Calculate the z component
	);
}

double Vector3D::inner(const Vector3D& v) const
{
	return x*v.x + y*v.y + z*v.z;
}

void TetMesh::initialize()
{
	for (auto& elem : elems) {
		auto tetElem = dynamic_cast<TetElement*>(elem.get());
		for (int i = 0; i < 3; ++i) {
			tetElem->center[i] = 0;
			for (int j = 0; j < 4; ++j) {
				tetElem->center[i] += nodes[tetElem->nodes_id[j]].coords[i];
			}
			tetElem->center[i] /= 4;
		}
	}
}

// nothing is below other than readVTK
void Mesh::readVTK(const std::string& filename) {
    std::ifstream file(filename);
    if (!file.is_open()) {
        std::cerr << "Error opening file: " << filename << std::endl;
        return;
    }

    std::string line;
    while (std::getline(file, line)) {
        std::istringstream iss(line);
        std::string keyword;
        iss >> keyword;

        if (keyword == "POINTS") {
            size_t numPoints;
            std::string dataType;
            iss >> numPoints >> dataType;
            nodes.resize(numPoints);
            for (size_t i = 0; i < numPoints; ++i) {
                file >> nodes[i].coords[0] >> nodes[i].coords[1] >> nodes[i].coords[2];
            }
        }
        else if (keyword == "CELLS") {
            size_t numCells, totalSize;
            iss >> numCells >> totalSize;
            for (size_t i = 0; i < numCells; ++i) {
                size_t cellSize;
                file >> cellSize;
                if (cellSize == 8) {
                    auto hexElem = std::make_unique<HexElement>();
                    for (size_t j = 0; j < 8; ++j) {
                        file >> hexElem->nodes_id[j];
                    }
                    elems.push_back(std::move(hexElem));
                }
                else if (cellSize == 4) {
                    auto tetElem = std::make_unique<TetElement>();
                    for (size_t j = 0; j < 4; ++j) {
                        file >> tetElem->nodes_id[j];
                    }
                    elems.push_back(std::move(tetElem));
                }
                else {
                    std::cerr << "Unsupported cell size: " << cellSize << std::endl;
                }
            }
        }
        else if (keyword == "POINT_DATA") {
            size_t numPoints;
            iss >> numPoints;
            
            // 读取下一行的数据类型信息
            std::getline(file, line);
            std::istringstream dataTypeLine(line);
            std::string dataTypeKeyword, dataName, dataType;
            size_t numComp;
            dataTypeLine >> dataTypeKeyword >> dataName >> dataType >> numComp;

            if (dataTypeKeyword == "SCALARS") {
                // 读取 LOOKUP_TABLE 行
                std::getline(file, line);
                
                // 读取每个点的类型数据
                for (size_t i = 0; i < numPoints; ++i) {
                    file >> nodes[i].pointType;
                }
            }
        }
        else if (keyword == "CELL_DATA") {
            size_t numCells;
            iss >> numCells;
            // Read the next line for the data type
            while(std::getline(file, line)) {
                std::istringstream dataTypeLine(line);
                std::string dataTypeKeyword, dataName, dataType;
                size_t numComp;
                dataTypeLine >> dataTypeKeyword >> dataName >> dataType >> numComp;

                if (dataTypeKeyword == "SCALARS") {
                    std::getline(file, line); // Read LOOKUP_TABLE line
                    if (dataName == "rocktype") {
                        // Read rocktype data
                        for (size_t i = 0; i < numCells; ++i) {
                            file >> elems[i]->rockType;
                        }
                    } else if (dataTypeKeyword == "SCALARS" && dataType == "float") {
                        for (size_t i = 0; i < numCells; ++i) {
                            double data;
                            file >> data;
                            elems[i]->data[dataName] = data;
                        }
                    }
                    else {
                        std::cerr << "Unsupported data type: " << dataType << std::endl;
                    }
                }
            }
        }
    }

    file.close();
}